import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import json
import h5py
from tqdm import tqdm
import pmcx
import torch

class MCXSimulation:
    def __init__(self, config_path=None):
        """
        初始化MCX仿真类
        
        Args:
            config_path: MCX配置文件路径
        """
        self.config = self._load_config(config_path) if config_path else {}
        self.results = {}
        self.device = torch.device('cpu')  # 强制使用CPU
        
    def _load_config(self, config_path):
        """加载MCX配置文件"""
        with open(config_path, 'r') as f:
            return json.load(f)
    
    def set_medium_properties(self, medium_type, properties):
        """
        设置介质属性
        
        Args:
            medium_type: 介质类型（如：'tissue', 'water'等）
            properties: 介质属性字典（包含mu_a, mu_s, g等）
        """
        if 'medium' not in self.config:
            self.config['medium'] = {}
        self.config['medium'][medium_type] = properties
    
    def set_source_properties(self, source_type, properties):
        """
        设置光源属性
        
        Args:
            source_type: 光源类型（如：'point', 'gaussian'等）
            properties: 光源属性字典
        """
        if 'source' not in self.config:
            self.config['source'] = {}
        self.config['source'][source_type] = properties
    
    def run_simulation(self, output_path=None):
        """
        运行MCX仿真
        
        Args:
            output_path: 输出文件路径
        """
        # 准备PMCX输入参数
        cfg = {
            'dim': self.config['simulation']['dim'],
            'unitinmm': self.config['simulation']['unitinmm'],
            'tstart': self.config['simulation']['tstart'],
            'tend': self.config['simulation']['tend'],
            'tstep': self.config['simulation']['tstep'],
            'nphoton': self.config['simulation']['nphoton'],
            'seed': self.config['simulation']['seed'],
            'vol': np.ones(self.config['simulation']['dim'][:3]),  # 创建默认体积
            'srcpos': np.array([50, 50, 0]),  # 光源位置
            'srcdir': np.array([0, 0, 1]),    # 光源方向
            'srcparam1': np.array([1, 0, 0, 0]),  # 光源参数
            'srcparam2': np.array([0, 0, 0, 0]),  # 光源参数
            'srcpattern': np.array([1]),      # 光源模式
            'prop': np.array([  # 默认介质属性
                [0, 0, 1, 1],    # 空气
                [0.1, 10.0, 0.9, 1.37]  # 组织
            ]),
            'gpu': 0  # 使用CPU
        }
        
        # 运行仿真
        self.results = pmcx.run(cfg)
        
        # 保存结果
        if output_path:
            self.save_results(output_path)
    
    def save_results(self, output_path):
        """
        保存仿真结果
        
        Args:
            output_path: 输出文件路径
        """
        with h5py.File(output_path, 'w') as f:
            for key, value in self.results.items():
                if isinstance(value, torch.Tensor):
                    f.create_dataset(key, data=value.cpu().numpy())
                else:
                    if isinstance(value, list):
                        f.create_dataset(key, data=value)
                    else:
                        print(value)
                        print(11)
    
    def load_results(self, result_path):
        """
        加载仿真结果
        
        Args:
            result_path: 结果文件路径
        """
        with h5py.File(result_path, 'r') as f:
            self.results = {key: torch.tensor(f[key][:], device=self.device) 
                          for key in f.keys()}
    
    def visualize_results(self, plot_type='fluence'):
        """
        可视化仿真结果
        
        Args:
            plot_type: 可视化类型（'fluence', 'photon_path'等）
        """
        if not self.results:
            print("没有可用的仿真结果")
            return
        
        if plot_type == 'fluence':
            plt.figure(figsize=(10, 8))
            fluence = self.results['fluence'].cpu().numpy()
            plt.imshow(fluence, cmap='jet')
            plt.colorbar(label='Fluence (J/cm²)')
            plt.title('Photon Fluence Distribution')
            plt.show()
        elif plot_type == 'photon_path':
            # 可视化光子路径
            paths = self.results['photon_path'].cpu().numpy()
            plt.figure(figsize=(10, 8))
            for path in paths[:100]:  # 只显示前100条路径
                plt.plot(path[:, 0], path[:, 2], 'b-', alpha=0.1)
            plt.title('Photon Paths')
            plt.show()
        elif plot_type == 'flux':
            # 可视化光子路径
            plt.figure(figsize=(10, 8))
            fluence = self.results['flux']
            plt.imshow(fluence, cmap='jet')
            plt.colorbar(label='Fluence (J/cm²)')
            plt.title('Photon Fluence Distribution')
            plt.show()
        elif plot_type == 'stat':
            # 可视化光子路径
            paths = self.results['stat']
            plt.figure(figsize=(10, 8))
            for path in paths[:300]:  # 只显示前100条路径
                plt.plot(path[:, 0], path[:, 2], 'b-', alpha=0.1)
            plt.title('Photon Paths')
            plt.show()

def main():
    # 示例使用
    sim = MCXSimulation('data/config.json')
    
    # 运行仿真
    sim.run_simulation('results/simulation_result.h5')
    
    # 可视化结果
    sim.visualize_results('flux')  # flux fluence
    sim.visualize_results('stat') # stat photon_path

if __name__ == "__main__":
    main() 